import argparse
import re

def convert_prototxt(input_file, output_file, width, height, classes):
    with open(input_file, 'r') as f:
        content = f.read()
    
    # Replace input shape parameters
    content = re.sub(r'input_dim: \d+\ninput_dim: \d+\ninput_dim: \d+\ninput_dim: \d+', 
                    f'input_shape {{\n  dim: 1\n  dim: 3\n  dim: {height}\n  dim: {width}\n}}', content)
    
    # Replace class numbers
    if classes != 80:  # Default is 80 classes
        # Update classes in yolo parameters
        content = re.sub(r'classes: 80', f'classes: {classes}', content)
        
        # Update output dimensions in the final conv layers (3 detection heads)
        # Each output has (4 + 1 + classes) * 3 anchors = (5 + classes) * 3
        new_output = (5 + classes) * 3
        content = re.sub(r'num_output: 255', f'num_output: {new_output}', content)
    
    # Fix the upsample layer parameters
    content = re.sub(r'scale: 2', 'scale: 1\n        stride: 2', content)
    
    # Add img_info input
    if 'input: "img_info"' not in content:
        content = content.replace('input: "data"', 'input: "data"\ninput: "img_info"')
        # Add img_info shape after the data input shape
        content = re.sub(r'input_shape \{\n  dim: 1\n  dim: 3\n  dim: \d+\n  dim: \d+\n\}',
                        r'\g<0>\n\ninput_shape {\n  dim: 1\n  dim: 4\n}', content)
    
    # Add detection output layer
    if 'name: "detection_out3"' not in content:
        yolo_output_layer = """
layer {
name: "detection_out3"
type: "YoloV3DetectionOutput"
bottom: "yolo1_coords" 
bottom: "yolo2_coords"
bottom: "yolo3_coords"
bottom: "yolo1_obj"
bottom: "yolo2_obj" 
bottom: "yolo3_obj"
bottom: "yolo1_classes"
bottom: "yolo2_classes"
bottom: "yolo3_classes" 
bottom: "img_info"
top: "box_out"
top: "box_out_num"
yolov3_detection_output_param {
  boxes: 3
  classes: %d
  relative: true
  obj_threshold: 0.5
  score_threshold: 0.5
  iou_threshold: 0.45
  pre_nms_topn: 512
  post_nms_topn: 1024
  biases_high: 10
  biases_high: 13
  biases_high: 16
  biases_high: 30
  biases_high: 33
  biases_high: 23
  biases_mid: 30
  biases_mid: 61  
  biases_mid: 62
  biases_mid: 45
  biases_mid: 59
  biases_mid: 119
  biases_low: 116
  biases_low: 90  
  biases_low: 156
  biases_low: 198
  biases_low: 373
  biases_low: 326
  }
}""" % classes
        
        # Add Yolo layers before the detection output
        yolo_layer_template = """
layer {
	bottom: "layer%s-conv"
	top: "yolo%s_coords"
	top: "yolo%s_obj"
	top: "yolo%s_classes"
	name: "yolo%s"
	type: "Yolo"
	yolo_param {
		boxes: 3
		coords: 4
		classes: %d  
		yolo_version: "V3"
		softmax: true
		background: false
	}
}"""
        
        # Find the last convolution layer (layer82-conv, layer94-conv, layer106-conv)
        yolo_patterns = [
            ('82', '1'), # First detection head
            ('94', '2'), # Second detection head
            ('106', '3') # Third detection head
        ]
        
        for layer_num, yolo_num in yolo_patterns:
            # Create the yolo layer using the template
            yolo_layer = yolo_layer_template % (layer_num, yolo_num, yolo_num, yolo_num, yolo_num, classes)
            
            # Pattern to find the end of the conv layer
            pattern = f"layer {{\n    bottom: \"layer{int(layer_num)-1}-conv\"\n    top: \"layer{layer_num}-conv\"\n(.+?)bias_term: true\n    }}\n}}"
            
            if f"name: \"yolo{yolo_num}\"" not in content:
                # Insert the yolo layer after the convolution layer
                match = re.search(pattern, content, re.DOTALL)
                if match:
                    # Find the position after the matched pattern
                    end_pos = match.end()
                    # Insert the yolo layer
                    content = content[:end_pos] + yolo_layer + content[end_pos:]
        
        # Add the detection output layer at the end
        content += yolo_output_layer
    
    with open(output_file, 'w') as f:
        f.write(content)

def main():
    parser = argparse.ArgumentParser(description='Convert Darknet to Caffe prototxt with configurable parameters')
    parser.add_argument('input', help='Input prototxt file')
    parser.add_argument('output', help='Output prototxt file')
    parser.add_argument('--width', type=int, default=416, help='Input width (default: 416)')
    parser.add_argument('--height', type=int, default=416, help='Input height (default: 416)')
    parser.add_argument('--classes', type=int, default=80, help='Number of classes (default: 80)')
    
    args = parser.parse_args()
    
    convert_prototxt(args.input, args.output, args.width, args.height, args.classes)
    print(f"Converted {args.input} to {args.output} with dimensions {args.width}x{args.height} and {args.classes} classes")

if __name__ == "__main__":
    main()

# 使用实例 python convert_prototxt.py input.prototxt output.prototxt --width xxx --height xxx --classes xxx